# 指定文件编码为UTF-8
# coding: utf-8

"""
optimizer.py 功能说明：
1. 实现多种神经网络优化算法
2. 包含6种常用优化器：
   - SGD: 随机梯度下降
   - Momentum: 动量梯度下降
   - Nesterov: Nesterov加速梯度
   - AdaGrad: 自适应学习率
   - RMSprop: 均方根传播
   - Adam: 自适应矩估计
3. 所有优化器统一接口，方便替换使用
4. 支持参数自动更新和状态维护

关键特性：
- 每种优化器都实现update()方法
- 支持参数分组更新
- 包含学习率衰减机制
- 实现梯度一阶和二阶矩估计
- 数值稳定性处理(如添加1e-7小常数)
"""

import numpy as np

class SGD:
    """随机梯度下降法(Stochastic Gradient Descent)

    参数:
        lr: 学习率(默认0.01)
    特点:
        - 简单易实现
        - 需要手动调整学习率
        - 可能收敛慢且不稳定
    """
    def __init__(self, lr=0.01):
        self.lr = lr  # 学习率

    def update(self, params, grads):
        """参数更新"""
        for key in params.keys():
            params[key] -= self.lr * grads[key]  # 简单梯度下降


class Momentum:
    """动量梯度下降法(Momentum SGD)

    参数:
        lr: 学习率(默认0.01)
        momentum: 动量系数(默认0.9)
    特点:
        - 加入动量项加速收敛
        - 减少震荡
        - 有助于逃离局部极小值
    """
    def __init__(self, lr=0.01, momentum=0.9):
        self.lr = lr
        self.momentum = momentum  # 动量系数
        self.v = None  # 速度

    def update(self, params, grads):
        """参数更新"""
        if self.v is None:  # 初始化速度
            self.v = {}
            for key, val in params.items():
                self.v[key] = np.zeros_like(val)

        for key in params.keys():
            # 更新速度(动量项+梯度项)
            self.v[key] = self.momentum*self.v[key] - self.lr*grads[key]
            params[key] += self.v[key]  # 更新参数


class Nesterov:
    """Nesterov加速梯度法

    参数:
        lr: 学习率(默认0.01)
        momentum: 动量系数(默认0.9)
    特点:
        - Momentum的改进版
        - 在计算梯度时考虑动量
        - 收敛速度更快
    """
    def __init__(self, lr=0.01, momentum=0.9):
        self.lr = lr
        self.momentum = momentum
        self.v = None

    def update(self, params, grads):
        """参数更新"""
        if self.v is None:
            self.v = {}
            for key, val in params.items():
                self.v[key] = np.zeros_like(val)

        for key in params.keys():
            self.v[key] *= self.momentum
            self.v[key] -= self.lr * grads[key]
            # Nesterov双重动量更新
            params[key] += self.momentum * self.momentum * self.v[key]
            params[key] -= (1 + self.momentum) * self.lr * grads[key]


class AdaGrad:
    """自适应梯度法(AdaGrad)

    参数:
        lr: 学习率(默认0.01)
    特点:
        - 自适应调整学习率
        - 适合稀疏数据
        - 学习率会单调递减
    """
    def __init__(self, lr=0.01):
        self.lr = lr
        self.h = None  # 梯度平方累积

    def update(self, params, grads):
        """参数更新"""
        if self.h is None:
            self.h = {}
            for key, val in params.items():
                self.h[key] = np.zeros_like(val)

        for key in params.keys():
            self.h[key] += grads[key] * grads[key]  # 累积梯度平方
            # 自适应学习率更新
            params[key] -= self.lr * grads[key] / (np.sqrt(self.h[key]) + 1e-7)


class RMSprop:
    """均方根传播法(RMSprop)

    参数:
        lr: 学习率(默认0.01)
        decay_rate: 衰减率(默认0.99)
    特点:
        - AdaGrad的改进版
        - 引入衰减系数解决学习率单调递减问题
        - 适合非平稳目标
    """
    def __init__(self, lr=0.01, decay_rate=0.99):
        self.lr = lr
        self.decay_rate = decay_rate  # 衰减率
        self.h = None

    def update(self, params, grads):
        """参数更新"""
        if self.h is None:
            self.h = {}
            for key, val in params.items():
                self.h[key] = np.zeros_like(val)

        for key in params.keys():
            # 指数移动平均更新梯度平方
            self.h[key] *= self.decay_rate
            self.h[key] += (1 - self.decay_rate) * grads[key] * grads[key]
            params[key] -= self.lr * grads[key] / (np.sqrt(self.h[key]) + 1e-7)


class Adam:
    """自适应矩估计(Adam)

    参数:
        lr: 学习率(默认0.001)
        beta1: 一阶矩衰减率(默认0.9)
        beta2: 二阶矩衰减率(默认0.999)
    特点:
        - 结合Momentum和RMSprop
        - 自适应调整学习率
        - 偏置校正
        - 通常效果最好
    """
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999):
        self.lr = lr
        self.beta1 = beta1  # 一阶矩衰减率
        self.beta2 = beta2  # 二阶矩衰减率
        self.iter = 0  # 迭代次数
        self.m = None  # 一阶矩(梯度均值)
        self.v = None  # 二阶矩(梯度方差)

    def update(self, params, grads):
        """参数更新"""
        if self.m is None:
            self.m, self.v = {}, {}
            for key, val in params.items():
                self.m[key] = np.zeros_like(val)
                self.v[key] = np.zeros_like(val)

        self.iter += 1
        # 计算偏置校正后的学习率
        lr_t = self.lr * np.sqrt(1.0 - self.beta2**self.iter) / (1.0 - self.beta1**self.iter)

        for key in params.keys():
            # 更新一阶矩估计(动量)
            self.m[key] += (1 - self.beta1) * (grads[key] - self.m[key])
            # 更新二阶矩估计(自适应学习率)
            self.v[key] += (1 - self.beta2) * (grads[key]**2 - self.v[key])

            # 参数更新(带偏置校正)
            params[key] -= lr_t * self.m[key] / (np.sqrt(self.v[key]) + 1e-7)
